package com.example.picturebacked.manager.sharding;

import com.baomidou.mybatisplus.extension.toolkit.SqlRunner;
import com.example.picturebacked.model.entity.Space;
import com.example.picturebacked.model.enums.SpaceLevelEnum;
import com.example.picturebacked.model.enums.SpaceTypeEnum;
import com.example.picturebacked.service.SpaceService;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.mode.metadata.MetaDataContexts;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @Autheor LBT
 * 根据生成的表名动态生成分片键，并返在读取配置文件时获取分片数量
 */
//@Component
@Slf4j
public class DynamicShardingManager {

    @Resource
    private DataSource dataSource;

    @Resource
    private SpaceService spaceService;

    // 逻辑表名，用于动态生成分表名的前缀
    private static final String LOGIC_TABLE_NAME = "picture";

    // 设置配置文件中的数据库名称
    public static final String DATEBASE_NAME = "logic_db";

    @PostConstruct
    public void initialize() {
        log.info("初始化动态分表配置...");
        updateShardingTableNodes();
    }

    /**
     * 更新 ShardingSphere 的 actual_data_nodes 动态表名配置
     */
    private void updateShardingTableNodes() {
        Set<String> tableNames =  fetchAllPictureTableNames();
        // 获取动态生成的表名，将其拼接上数据库名的前缀,如下面所示，用于配置到配置文件的动态表名配置中
        // picture_backed.picture_123165746545, picture_backed.picture_123165746546
        String newActualDataNodes = tableNames.stream().map(tabelName -> "picture_backed." + tabelName)
                .collect(Collectors.joining(","));
        log.info("动态分表 actual-data-nodes 配置：{}",newActualDataNodes);

        ContextManager contextManager = getContextManager();

        ShardingSphereRuleMetaData ruleMetaData = contextManager.getMetaDataContexts()
                .getMetaData()
                .getDatabases()
                .get(DATEBASE_NAME)
                .getRuleMetaData();

        Optional<ShardingRule> shardingRule = ruleMetaData.findSingleRule(ShardingRule.class);
        if (shardingRule.isPresent()) {
            ShardingRuleConfiguration ruleConfig= (ShardingRuleConfiguration) shardingRule.get().getConfiguration();
            List<ShardingTableRuleConfiguration> updateRules = ruleConfig.getTables().stream().map(oldTableRule -> {
                if (LOGIC_TABLE_NAME.equals(oldTableRule.getLogicTable())) {
                    ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, newActualDataNodes);
                    newTableRuleConfig.setDatabaseShardingStrategy(oldTableRule.getDatabaseShardingStrategy());
                    newTableRuleConfig.setTableShardingStrategy(oldTableRule.getTableShardingStrategy());
                    newTableRuleConfig.setKeyGenerateStrategy(oldTableRule.getKeyGenerateStrategy());
                    newTableRuleConfig.setAuditStrategy(oldTableRule.getAuditStrategy());
                    return newTableRuleConfig;
                }
                return oldTableRule;
            }).collect(Collectors.toList());
            ruleConfig.setTables(updateRules);
            contextManager.alterRuleConfiguration(DATEBASE_NAME, Collections.singleton(ruleConfig));
            contextManager.reloadDatabase(DATEBASE_NAME);
            log.info("更新动态分表配置成功！");
        }else {
            log.error("未找到分片规则！动态分表更新失败");
        }
    }

    /**
     * 获取  ShardingSphere ContextMnager
     */
    private ContextManager getContextManager() {
        try(ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class)){
            return connection.getContextManager();
        }catch (SQLException e){
            throw new RuntimeException("获取ShardingSphere ContextManager 失败",e);
        }
    }

    /**
     * 获取所有动态表明，包括初始表picture和分表picture_{spaceId}
     * @return
     */
    private Set<String> fetchAllPictureTableNames() {
        Set<Long> spaceIds = spaceService.lambdaQuery()
                .eq(Space::getSpaceType, SpaceTypeEnum.TEAM.getValue())
                .list()
                .stream()
                .map(Space::getId)
                .collect(Collectors.toSet());
        Set<String> tableNames = spaceIds.stream()
                .map(spaceId -> LOGIC_TABLE_NAME + "_" + spaceId)
                .collect(Collectors.toSet());
        tableNames.add(LOGIC_TABLE_NAME);// 添加初始化逻辑表
        return tableNames;
    }


    /**
     * 动态创建分表方法，通过拼接SQL的方式创建出和Picture表结构一样的分表
     * @param space
     */
    public void  createSpacePictureTable(Space  space) {
        // 动态创建分表
        if (space.getSpaceType() == SpaceTypeEnum.TEAM.getValue() && space.getSpaceLevel() == SpaceLevelEnum.FLAGSHIP.getValue()) {
            Long id = space.getId();
            String tableName = "picture_" + id;
            // 创建新表
            String createTableSQL = "CREATE TABLE " + tableName + " LIKE picture";
            try{
                SqlRunner.db().update(createTableSQL);
                // 更新分表
                updateShardingTableNodes();
            }catch (Exception e){
                log.error("创建图片空间分表失败！空间id = {}",id);
            }
        }
    }

}
